[not for land] hook up MX to CUDA 12.8 cuBLAS MX gemm #1625

vkuzo · 2025-01-27T18:26:03Z

Summary:

Requires https://github.com/pytorch/pytorch/pull/145562/files

None of this is for land - just testing for now as we work on a long
term support plan.

Status:

performance

[done] torch._scaled_mm performance as expected: seeing about 2x speedup from bf16 to mxfp8, and 4x from bf16 to nvfp4
[not in this PR] fast scaling kernels

numerics - mxfp8

[done] hello world
[done] check numerics with A and B set to eye-one and scales set to all-ones
[done] numerics match fp32 baseline and mx emulation with a single tile (M == K == N == 128), with torch.randn inputs
[done] numerics match fp32 baseline for multiple tiles, of various shapes
[done] (not in this PR) train llama with real MX gemm and get decent numerics

numerics - nvfp4

[done] hello world
[not done] repeat all steps of mxfp8 numerics testing

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:

[ghstack-poisoned]

vkuzo · 2025-01-27T18:26:04Z

Stack from ghstack (oldest at bottom):

pytorch-bot · 2025-01-27T18:26:06Z

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/1625

📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 9 New Failures

As of commit c8d768e with merge base 47f96f1 ():

NEW FAILURES - The following jobs have failed:

Code Analysis with Ruff / build (3.9) (gh)
Run Regression Tests / test (CPU 2.3, linux.4xlarge, torch==2.3.0 --index-url https://download.pytorch.org/whl/cpu, cpu) / linux-job (gh)
RuntimeError: Command docker exec -t 7f882dc6bdea44aebef76ec502652f71705f531d9cd9e88ee21c4b172fd5b3d7 /exec failed with exit code 2
Run Regression Tests / test (CPU 2.4, linux.4xlarge, torch==2.4.0 --index-url https://download.pytorch.org/whl/cpu, cpu) / linux-job (gh)
RuntimeError: Command docker exec -t 75b4f324e1450773eca8fc1ecc615fcebe886d69a9d730d0bede86073ee66a64 /exec failed with exit code 2
Run Regression Tests / test (CPU 2.5.1, linux.4xlarge, torch==2.5.1 --index-url https://download.pytorch.org/whl/cpu, cpu) / linux-job (gh)
RuntimeError: Command docker exec -t 54a8ce1343fbbd998a0901d5a48eb4f00c09f0448677eac433d791d8fc9d450c /exec failed with exit code 2
Run Regression Tests / test (CUDA 2.3, linux.g5.12xlarge.nvidia.gpu, torch==2.3.0, cuda, 12.1) / linux-job (gh)
RuntimeError: Command docker exec -t 1ff80f76f1b5372b57ba5ea5bca2db252473d7d6932ca80f43e03083b9749b7d /exec failed with exit code 2
Run Regression Tests / test (CUDA 2.4, linux.g5.12xlarge.nvidia.gpu, torch==2.4.0, cuda, 12.1) / linux-job (gh)
RuntimeError: Command docker exec -t 3b9b5bc189e93b66b274daa35332ecdff8c1067f493c318efc5d94b97ea05752 /exec failed with exit code 2
Run Regression Tests / test (CUDA 2.5.1, linux.g5.12xlarge.nvidia.gpu, torch==2.5.1 --index-url https://download.pytorch... / linux-job (gh)
RuntimeError: Command docker exec -t a6fe078e551843efca80082a93dedba34635440c451208b06338a60a4b359ff2 /exec failed with exit code 2
Run Regression Tests / test-nightly (CPU Nightly, linux.4xlarge, --pre torch==2.7.0.dev20250122 --index-url https://down... / linux-job (gh)
RuntimeError: Command docker exec -t 37e9a2e0e4faaa4faff07949a20d9d057ad77a7b4e69a7126a60beb324815415 /exec failed with exit code 2
Run Regression Tests / test-nightly (CUDA Nightly, linux.g5.12xlarge.nvidia.gpu, --pre torch==2.7.0.dev20250122 --index-... / linux-job (gh)
RuntimeError: Command docker exec -t e3f78c4a3ff9cdcafbc114f0da32f5418b4de063ddd9a0ea495d4f195cbc4760 /exec failed with exit code 2

This comment was automatically generated by Dr. CI and updates every 15 minutes.

Summary: Requires https://github.com/pytorch/pytorch/pull/145562/files None of this is for land - just testing for now as we work on a long term support plan. Test Plan: Reviewers: Subscribers: Tasks: Tags: ghstack-source-id: dcde36f ghstack-comment-id: 2616580142 Pull Request resolved: #1625

[ghstack-poisoned]

Summary: Requires https://github.com/pytorch/pytorch/pull/145562/files None of this is for land - just testing for now as we work on a long term support plan. Test Plan: Reviewers: Subscribers: Tasks: Tags: ghstack-source-id: 9d7e307 ghstack-comment-id: 2616580142 Pull Request resolved: #1625

[ghstack-poisoned]

Summary: Requires https://github.com/pytorch/pytorch/pull/145562/files None of this is for land - just testing for now as we work on a long term support plan. Test Plan: Reviewers: Subscribers: Tasks: Tags: ghstack-source-id: 5442d55 ghstack-comment-id: 2616580142 Pull Request resolved: #1625

[ghstack-poisoned]

Summary: Requires https://github.com/pytorch/pytorch/pull/145562/files None of this is for land - just testing for now as we work on a long term support plan. Test Plan: Reviewers: Subscribers: Tasks: Tags: ghstack-source-id: cd9d044 ghstack-comment-id: 2616580142 Pull Request resolved: #1625

[ghstack-poisoned]

Summary: Requires https://github.com/pytorch/pytorch/pull/145562/files None of this is for land - just testing for now as we work on a long term support plan. Test Plan: Reviewers: Subscribers: Tasks: Tags: ghstack-source-id: bebf96f ghstack-comment-id: 2616580142 Pull Request resolved: #1625

[ghstack-poisoned]

Summary: Requires https://github.com/pytorch/pytorch/pull/145562/files None of this is for land - just testing for now as we work on a long term support plan. Test Plan: Reviewers: Subscribers: Tasks: Tags: ghstack-source-id: af1626e ghstack-comment-id: 2616580142 Pull Request resolved: #1625

vkuzo · 2025-01-28T00:31:01Z

test/prototype/mx_formats/test_mx_linear.py

+    #       11
+    #       11
+    #
+    # for each row in rows of A:


vkuzo · 2025-01-28T00:31:10Z

test/prototype/mx_formats/test_mx_linear.py

+    #       11
+    #
+    # for each row in rows of A:
+    #   for each ol in cols of A:


[ghstack-poisoned]

Summary: Requires https://github.com/pytorch/pytorch/pull/145562/files None of this is for land - just testing for now as we work on a long term support plan. Test Plan: Reviewers: Subscribers: Tasks: Tags: ghstack-source-id: e7a7997 ghstack-comment-id: 2616580142 Pull Request resolved: #1625

[ghstack-poisoned]

Summary: Requires https://github.com/pytorch/pytorch/pull/145562/files None of this is for land - just testing for now as we work on a long term support plan. Test Plan: Reviewers: Subscribers: Tasks: Tags: ghstack-source-id: 946b39a ghstack-comment-id: 2616580142 Pull Request resolved: #1625

[ghstack-poisoned]

Summary: Requires https://github.com/pytorch/pytorch/pull/145562/files None of this is for land - just testing for now as we work on a long term support plan. Test Plan: Reviewers: Subscribers: Tasks: Tags: ghstack-source-id: e8b0cfb ghstack-comment-id: 2616580142 Pull Request resolved: #1625

[ghstack-poisoned]

Summary: Requires https://github.com/pytorch/pytorch/pull/145562/files None of this is for land - just testing for now as we work on a long term support plan. Test Plan: Reviewers: Subscribers: Tasks: Tags: ghstack-source-id: 8a7803b ghstack-comment-id: 2616580142 Pull Request resolved: #1625

[ghstack-poisoned]

Summary: Requires https://github.com/pytorch/pytorch/pull/145562/files None of this is for land - just testing for now as we work on a long term support plan. Test Plan: Reviewers: Subscribers: Tasks: Tags: ghstack-source-id: 599acd3 ghstack-comment-id: 2616580142 Pull Request resolved: #1625

[ghstack-poisoned]

Summary: Requires https://github.com/pytorch/pytorch/pull/145562/files None of this is for land - just testing for now as we work on a long term support plan. Test Plan: Reviewers: Subscribers: Tasks: Tags: ghstack-source-id: d3762db ghstack-comment-id: 2616580142 Pull Request resolved: #1625

[ghstack-poisoned]

Summary: Requires https://github.com/pytorch/pytorch/pull/145562/files None of this is for land - just testing for now as we work on a long term support plan. Test Plan: Reviewers: Subscribers: Tasks: Tags: ghstack-source-id: d24fb59 ghstack-comment-id: 2616580142 Pull Request resolved: #1625

Update

a6c9f1c

[ghstack-poisoned]

vkuzo mentioned this pull request Jan 27, 2025

skip failing MX tests on cuda capability 10.0 #1624

Merged

facebook-github-bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Jan 27, 2025

vkuzo added the topic: not user facing Use this tag if you don't want this PR to show up in release notes label Jan 27, 2025

Update

5966516

[ghstack-poisoned]

Update

759797b

[ghstack-poisoned]

Update

ee80fa7

[ghstack-poisoned]

Update

c04dbef

[ghstack-poisoned]

Update

90c6e43

[ghstack-poisoned]

vkuzo commented Jan 28, 2025

View reviewed changes

test/prototype/mx_formats/test_mx_linear.py

# 11

# 11

#

# for each row in rows of A:

Copy link

Contributor Author

vkuzo Jan 28, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

scale_A

vkuzo commented Jan 28, 2025

View reviewed changes

test/prototype/mx_formats/test_mx_linear.py

# 11

#

# for each row in rows of A:

# for each ol in cols of A:

Copy link

Contributor Author

vkuzo Jan 28, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

scale_A

Update

0a8ce19

[ghstack-poisoned]

Update

686be3d

[ghstack-poisoned]

Update

8ecba5f

[ghstack-poisoned]

Update

c5df6d8

[ghstack-poisoned]

Update

8daa0d0

[ghstack-poisoned]

Update

be7806c

[ghstack-poisoned]

Update

c8d768e

[ghstack-poisoned]

drisspg mentioned this pull request Jan 29, 2025

Add mx_fp8_bf16 kernel #1637

Merged

vkuzo closed this Feb 25, 2025

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[not for land] hook up MX to CUDA 12.8 cuBLAS MX gemm #1625

[not for land] hook up MX to CUDA 12.8 cuBLAS MX gemm #1625

vkuzo commented Jan 27, 2025 •

edited

Loading

vkuzo commented Jan 27, 2025 •

edited

Loading

pytorch-bot bot commented Jan 27, 2025 •

edited

Loading

vkuzo Jan 28, 2025

vkuzo Jan 28, 2025

[not for land] hook up MX to CUDA 12.8 cuBLAS MX gemm #1625

[not for land] hook up MX to CUDA 12.8 cuBLAS MX gemm #1625

Conversation

vkuzo commented Jan 27, 2025 • edited Loading

performance

numerics - mxfp8

numerics - nvfp4

vkuzo commented Jan 27, 2025 • edited Loading

pytorch-bot bot commented Jan 27, 2025 • edited Loading

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/1625

❌ 9 New Failures

vkuzo Jan 28, 2025

Choose a reason for hiding this comment

vkuzo Jan 28, 2025

Choose a reason for hiding this comment

vkuzo commented Jan 27, 2025 •

edited

Loading

vkuzo commented Jan 27, 2025 •

edited

Loading

pytorch-bot bot commented Jan 27, 2025 •

edited

Loading